package main

import (
	"encoding/json"
	"net/http"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type WsConnect struct {
	ChIn  chan []byte
	ChOut chan []byte
}

type WsClientConn struct {
	Lock    *sync.Mutex
	Apps    map[uint]map[*websocket.Conn]WsConnect
	Devices map[uint]map[*websocket.Conn]WsConnect
	Users   map[uint]map[*websocket.Conn]WsConnect
}

var wsclient_pool WsClientConn

func (m *WsClientConn) Initialize() {
	m.Lock = new(sync.Mutex)
	m.Apps = make(map[uint]map[*websocket.Conn]WsConnect)
	m.Devices = make(map[uint]map[*websocket.Conn]WsConnect)
	m.Users = make(map[uint]map[*websocket.Conn]WsConnect)
}

func (m *WsClientConn) AddAppsConn(appdbid uint, conn *websocket.Conn, channel WsConnect) {
	m.Lock.Lock()
	defer m.Lock.Unlock()

	conns, ok := m.Apps[appdbid]
	if !ok {
		//empty, init.
		conns = make(map[*websocket.Conn]WsConnect)
		m.Apps[appdbid] = conns
	}
	conns[conn] = channel
}

func (m *WsClientConn) RemoveAppsConn(appdbid uint, conn *websocket.Conn) {
	m.Lock.Lock()
	defer m.Lock.Unlock()
	conns, ok := m.Apps[appdbid]
	if !ok {
		logger.Errorf("remove container apps by appid %d failed.", appdbid)
		return
	}
	_, ok = conns[conn]
	if !ok {
		logger.Errorf("remove container apps connect by appid %d failed, not exist.", appdbid)
		return
	}

	delete(conns, conn)
}

func (m *WsClientConn) AddUsersConn(userdbid uint, conn *websocket.Conn, channel WsConnect) {
	m.Lock.Lock()
	defer m.Lock.Unlock()

	conns, ok := m.Users[userdbid]
	if !ok {
		//empty, init.
		conns = make(map[*websocket.Conn]WsConnect)
		m.Users[userdbid] = conns
	}
	conns[conn] = channel
}

func (m *WsClientConn) RemoveUsersConn(userdbid uint, conn *websocket.Conn) {
	m.Lock.Lock()
	defer m.Lock.Unlock()
	conns, ok := m.Users[userdbid]
	if !ok {
		logger.Errorf("remove container Users by userdbid %d failed.", userdbid)
		return
	}
	_, ok = conns[conn]
	if !ok {
		logger.Errorf("remove container Users connect by userdbid %d failed, not exist.", userdbid)
		return
	}

	delete(conns, conn)
}

func (m *WsClientConn) appBoardcast(appdbid uint, msg []byte) {
	conns, ok := m.Apps[appdbid]
	if !ok {
		logger.Errorf("AppBoardcast apps by appid %d failed.", appdbid)
		return
	}

	for _, channel := range conns {
		channel.ChOut <- msg
	}
}

func (m *WsClientConn) userBoardcast(userdbid uint, msg []byte) {

	conns, ok := m.Users[userdbid]
	if !ok {
		logger.Errorf("user boardcast users by userdbid failed. %d", userdbid)
		return
	}

	for _, channel := range conns {
		channel.ChOut <- msg
	}
}

func (m *WsClientConn) Boardcast(device Device, msg []byte) {
	userdbid, err := QueryUserIDFromAppID(device.AppID)
	if err != nil {
		logger.Warningf("query user from appid %d failed", device.AppID)
		return
	}

	m.Lock.Lock()
	defer m.Lock.Unlock()

	m.appBoardcast(device.AppID, msg)
	m.userBoardcast(userdbid, msg)
}

func RpcDeviceRegister(device Device) {
	var jsondata struct {
		Method     string `json:"method"`
		DeviceDbID uint   `json:"device_id"`
		ChipID     string `json:"imei"`
		AppID      uint   `json:"app_id"`
	}

	jsondata.Method = "REGISTER"
	jsondata.DeviceDbID = device.ID
	jsondata.ChipID = device.ChipId
	jsondata.AppID = device.AppID

	data, err := json.Marshal(jsondata)
	if err != nil {
		logger.Warningf("json marshal failed.")
		return
	}
	wsclient_pool.Boardcast(device, data)
}

func RpcDataUploaded(device Device, data []byte) {
	var jsondata struct {
		Method     string `json:"method"`
		DeviceDbID uint   `json:"device_id"`
		ChipID     string `json:"imei"`
		AppID      uint   `json:"app_id"`
		Data       string `json:"data"`
	}

	jsondata.Method = "UPLOAD"
	jsondata.DeviceDbID = device.ID
	jsondata.ChipID = device.ChipId
	jsondata.AppID = device.AppID
	jsondata.Data = string(data)

	data, err := json.Marshal(jsondata)
	if err != nil {
		logger.Warningf("json marshal failed.")
		return
	}
	wsclient_pool.Boardcast(device, data)
}

func RpcDataSendAck(device Device, msgdbid uint) {
	var jsondata struct {
		Method     string `json:"method"`
		DeviceDbID uint   `json:"device_id"`
		ChipID     string `json:"imei"`
		AppID      uint   `json:"app_id"`
		MsgID      uint   `json:"msg_id"`
	}

	jsondata.Method = "DATACK"
	jsondata.DeviceDbID = device.ID
	jsondata.ChipID = device.ChipId
	jsondata.AppID = device.AppID
	jsondata.MsgID = msgdbid

	data, err := json.Marshal(jsondata)
	if err != nil {
		logger.Warningf("json marshal failed.")
		return
	}
	wsclient_pool.Boardcast(device, data)
}

func WebsocketProcHandler(conn *websocket.Conn, ChIn chan []byte) {
	for {
		msgType, msg, err := conn.ReadMessage()
		if err != nil {
			logger.Warning("websocket conn read message failed.", err)
			close(ChIn)
			break
		}
		logger.Infof("recv %s, format %s", msg, msgType)
		ChIn <- msg
	}
}

func WebSocketApiHandler(c *gin.Context) {
	app := c.MustGet(JWT_API_KEY_APP).(App)
	logger.Infof("websocket connection established, appid: %d", app.ID)
	v3wsupgrader := websocket.Upgrader{}
	conn, err := v3wsupgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.Warning("upgrade to websocket failed.")
		c.JSON(http.StatusNotFound, gin.H{"msg": "only accept websocket connection."})
		c.Abort()
		return
	}
	logger.Infof("recv app websocket connection %d", app.ID)
	wsChan := WsConnect{make(chan []byte), make(chan []byte)}
	wsclient_pool.AddAppsConn(app.ID, conn, wsChan)

	defer func() {
		conn.Close()
		wsclient_pool.RemoveAppsConn(app.ID, conn)
	}()

	go WebsocketProcHandler(conn, wsChan.ChIn)

	disconnected := false
	for {
		if disconnected {
			break
		}

		select {
		case v, ok := <-wsChan.ChIn:
			if !ok {
				logger.Warning("websocket rpc connection disconnected.")
				disconnected = true
				break
			}
			logger.Infof("data in... %s", string(v))
		case c := <-wsChan.ChOut:
			conn.WriteMessage(websocket.TextMessage, c)
		}
	}
	logger.Info("ws handler completed.")
}

func WsApiInitialize() {
	wsclient_pool.Initialize()
}
